{-# LANGUAGE OverloadedStrings, MultiParamTypeClasses, FlexibleInstances, ForeignFunctionInterface #-}

module Fits
  ( FitsImage(..)
  , HDU(..)
  , Header
  , Record(..)
  , Keyword(..)
  , Value(..)
  , Comment
  , writeFits
  ) where

import Data.String  
import Data.ByteString.Internal (c2w)
import Data.Int

import   Data.Array.IO      ( IOUArray, getBounds, freeze )
import   Data.Array.Unboxed ( UArray, (!) )
import          Data.Binary ( Binary(..), put )
import      Data.Binary.Put ( runPut )

import Control.Applicative  ( (<$>) )

import GHC.Float (double2Int)

import qualified Data.ByteString.Lazy as B

type Header  = [Record]
type Comment = String

data Record  = Record Keyword Value Comment | END

data Keyword = Keyword String deriving (Show)

data Value   = VInt Int | VFloat Float | VDouble Double | VString String | VBool Bool                 

data HDU     = HDUImage Header B.ByteString

instance IsString B.ByteString where
  fromString s = B.pack $ c2w <$> s

instance Show Record where
    show END = "END"
    show (Record (Keyword k) v _) = k ++ " = " ++ show v

instance Show Value where
    show (VInt     v) = show v
    show (VFloat   v) = show v
    show (VDouble  v) = show v
    show (VString  v) = v
    show (VBool True) = "T"
    show (VBool    _) = "F"

instance Show HDU where
    show (HDUImage h _) = (show <$> h) `joinBy` "\n"

class Serialize a where
  serialize :: a -> B.ByteString

instance Serialize Header where
  serialize h = round2880 $ B.concat $ serialize <$> h

instance Serialize Record where
  serialize END = padding 80 "END"
  serialize (Record k v c) = padding 80 $ sk `B.append` ("= ") `B.append` sv `B.append` sc
    where
      sk = serialize k
      sv = serialize v
      sc = serialize c

instance Serialize Keyword where
  serialize (Keyword k) = padding 8 $ fromString k

instance Serialize Value where
  serialize = fromString . show

instance Serialize Comment where
  serialize [] = ""
  serialize c  = fromString $ " / " ++ c

class FitsImage i where
  serializeImage :: Header -> i -> IO HDU

instance FitsImage (IOUArray (Int,Int) Double) where
  serializeImage header img = do
        b <- getBounds img
        a <- freeze img :: IO (UArray (Int,Int) Double)

        let
          ((x,y),(w,h)) = b

          minHeader  = [ Record (Keyword "SIMPLE") (VBool True)   "file conforms to FITS standard"
                       , Record (Keyword "BITPIX") (VInt 16)      "number of bits per data pixel"
                       , Record (Keyword "NAXIS")  (VInt 2)       "number of data axes"
                       , Record (Keyword "NAXIS1") (VInt (w+1-x)) "length of data axis 1"
                       , Record (Keyword "NAXIS2") (VInt (h+1-y)) "length of data axis 2"
                       , Record (Keyword "BZERO")  (VInt (2^15))  "physical value = BZERO + array value"
--                       , END
                       ]

          dat   = runPut $ mapM_ put $ map (\x -> round' x - 2^15) $ [a ! (i,j) | j <- [y..h], i <- [x..w]]

        return (HDUImage (minHeader ++ header ++ [END]) dat)

writeFits :: FilePath -> HDU -> IO ()
writeFits file (HDUImage header raw) = B.writeFile file (serialize header `B.append` raw)

{-# RULES
"truncate/Double->Int16"
    truncate = fromIntegral . double2Int :: Double -> Int16
 #-}

{-# INLINE round' #-}
round' :: Double -> Int16
round' x = (truncate :: Double -> Int16) (c_round x)

foreign import ccall unsafe "math.h round"
    c_round :: Double -> Double

joinBy :: [String] -> String -> String
joinBy (x:[]) _ = x
joinBy (x:xs) d = x ++ d ++ (xs `joinBy` d)
joinBy []     _ = []


padding :: Int64 -> B.ByteString -> B.ByteString
padding n s | l  > n = error $ "length string bigger than " ++ show l
            | l == n = s
            | otherwise = s `B.append` (B.take (n - l) $ B.repeat 0x20) 
   where
     l = B.length s

round2880 :: B.ByteString -> B.ByteString
round2880 s = padding b s where
    n = B.length s

    r = n `mod` 2880

    b = if r == 0 
         then n
         else 2880 + n - r